#include <cuda.h>
#include <cuda_runtime.h>
#include <cmath>
#include <cstdio>

// CUDA 错误检查宏
#define CUDA_CHECK_ERROR(call)                                         \
    do                                                                 \
    {                                                                  \
        cudaError_t err = call;                                        \
        if (err != cudaSuccess)                                        \
        {                                                              \
            fprintf(stderr, "CUDA Error: %s (err_num=%d) at %s:%d\n",  \
                    cudaGetErrorString(err), err, __FILE__, __LINE__); \
            exit(err);                                                 \
        }                                                              \
    } while (0)

// CUDA 内核：计算 Softmax
__global__ void softmax_kernel(const float *__restrict__ dev_input,
                               float *__restrict__ dev_out,
                               int rows, int cols)
{
    extern __shared__ float shared_mem[];

    // 每个块处理一行
    int row = blockIdx.x;
    if (row >= rows)
        return;

    // 计算当前行在输入和输出数组中的起始位置
    const float *input_row = dev_input + row * cols;
    float *output_row = dev_out + row * cols;

    // 指定线程在块内的索引
    int tid = threadIdx.x;
    int block_size = blockDim.x;

    // Step 1: 计算每行的最大值（用于数值稳定性）
    float max_val = -FLT_MAX;
    for (int i = tid; i < cols; i += block_size)
    {
        float val = input_row[i];
        if (val > max_val)
        {
            max_val = val;
        }
    }

    // 在共享内存中存储每个线程的局部最大值
    shared_mem[tid] = max_val;
    __syncthreads();

    // 使用并行归约计算全局最大值
    for (int stride = block_size / 2; stride > 0; stride >>= 1)
    {
        if (tid < stride)
        {
            if (shared_mem[tid + stride] > shared_mem[tid])
            {
                shared_mem[tid] = shared_mem[tid + stride];
            }
        }
        __syncthreads();
    }

    float row_max = shared_mem[0];

    // Step 2: 计算 exp(x[j] - max) 并累加
    float sum = 0.0f;
    for (int i = tid; i < cols; i += block_size)
    {
        float exp_val = expf(input_row[i] - row_max);
        shared_mem[tid] = exp_val; // 存储中间值用于后续归一化
        sum += exp_val;
    }

    // 在共享内存中存储每个线程的局部和
    shared_mem[tid] = sum;
    __syncthreads();

    // 使用并行归约计算全局和
    for (int stride = block_size / 2; stride > 0; stride >>= 1)
    {
        if (tid < stride)
        {
            shared_mem[tid] += shared_mem[tid + stride];
        }
        __syncthreads();
    }

    float row_sum = shared_mem[0];

    // Step 3: 归一化输出
    for (int i = tid; i < cols; i += block_size)
    {
        output_row[i] = shared_mem[tid] > 0 ? shared_mem[tid] / row_sum : 0.0f;
    }
}
